#include <math.h>
#include <stdlib.h>
#include "../GridException.h"
#include "helpers.h"
#include "../params.h"

#define CPU_USE_DOUBLE

static double* qforceX = NULL;
static double* qforceY = NULL;
static double* qforceZ = NULL;

static double (*computeLogPsi)(double r) = NULL;
static double (*computeLogPsiPrime)(double r) = NULL;
static double (*computeLogPsiDoublePrime)(double r) = NULL;

bool GridCPUInit()
{
  switch(getWaveFunction())
  {
    case kWaveFunctionDefault:
      computeLogPsi = defaultPsi;
      computeLogPsiPrime = defaultPsiPrime;
      computeLogPsiDoublePrime = defaultPsiDoublePrime;
      break;
    case kWaveFunctionMcMillan:
      computeLogPsi = mcMillanPsi;
      computeLogPsiPrime = mcMillanPsiPrime;
      computeLogPsiDoublePrime = mcMillanPsiDoublePrime;
      break;
    default:
      throw GridException("GridCPUInit():Invalid wave function");
  }

  return true;
}

void GridCPUComputeTotalEnergy(int deviceNum,
                               float* x,
                               float* y,
                               float* z,
                               int numParticles,
                               float* wavefunction,
                               float* potential,
                               float* kinetic)
{

  double dx;
  double dy;
  double dz;
  double dr;
  double pot = 0.0;
  double psi = 0.0;
  double gradientDivR = 0.0;
  double laplacian = 0.0;
  double kineticAccumulator = 0.0;

  //Init qforce vectors if needed
  if(qforceX == NULL)
  {
    qforceX = new double[numParticles];
    qforceY = new double[numParticles];
    qforceZ = new double[numParticles];
  }

  //Init energies
  *potential = 0.0f;
  *wavefunction = 0.0f;
  *kinetic = 0.0f;

  //Zero the quantum force vectors
  for(int i = 0; i < numParticles; i++)
  {
    qforceX[i] = 0.0;
    qforceY[i] = 0.0;
    qforceZ[i] = 0.0;
  }

  for(int i = 0; i < numParticles; i++)
  {
    for(int j = 0; j < i; j++)
    {
      dx = x[i] - x[j];
      dy = y[i] - y[j];
      dz = z[i] - z[j];

      dr = sqrt(dx * dx + dy * dy + dz * dz);

      pot += computePotential(dr);
      psi += computeLogPsi(dr);

      gradientDivR = computeLogPsiPrime(dr) / dr;
      laplacian = computeLogPsiDoublePrime(dr);

      qforceX[i] += dx * gradientDivR;
      qforceY[i] += dy * gradientDivR;
      qforceZ[i] += dz * gradientDivR;
      qforceX[j] -= dx * gradientDivR;
      qforceY[j] -= dy * gradientDivR;
      qforceZ[j] -= dz * gradientDivR;

      kineticAccumulator -= laplacian + 2.0 * gradientDivR;

    }
  }

  *wavefunction = psi;
  *potential = pot;

  //Add the quantum forces to the kinetic energy
  kineticAccumulator *= 2.0;

  for(int i = 0; i < numParticles; i++)
  {
    kineticAccumulator -=  qforceX[i] * qforceX[i] 
      + qforceY[i] * qforceY[i] 
      + qforceZ[i] * qforceZ[i];
  }

  kineticAccumulator *= 0.5 / gParticleMass;

  *kinetic = (float)kineticAccumulator;
}

bool GridCPUDeinit()
{
  if(qforceX != NULL)
  {
    delete[] qforceX;
  }
  if(qforceY != NULL)
  {
    delete[] qforceY;
  }
  if(qforceX != NULL)
  {
    delete[] qforceZ;
  }
  return true;
}
